package clear.experiment;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Random;

import clear.dep.DepNode;
import clear.dep.DepTree;
import clear.reader.DepReader;

import com.carrotsearch.hppc.IntOpenHashSet;
import com.carrotsearch.hppc.ObjectIntOpenHashMap;

public class DepGMCluster
{
	final int    MAX_ITER = 20;
	final double SQRT_2PI = Math.sqrt(2 * Math.PI);
	final double STOP     = 0.01;
	
	ObjectIntOpenHashMap<String> m_lexica;
	int[][]    a_vectors;	// i_N x *
	double[][] d_m;			// i_D x i_K
	double[]   d_s;			// i_K
	double[]   p_k;			// i_K
	double[][] p_kn;		// i_N x i_K
	double[]   d_mScalars;	// i_K	
	int        i_K, i_D, i_N;
	
	public DepGMCluster(String trainFile, int K)
	{
		init(trainFile, K);
		cluster();
	}
	
	public void init(String trainFile, int K)
	{
		i_K = K;
		
		initLexica (trainFile);
		initVectors(trainFile);
		initDistributions();
	}
	
	private void initLexica(String trainFile)
	{
		DepReader reader = new DepReader(trainFile, true);
		DepTree   tree;
		int       d = 1;
		
		System.out.print("Initializing lexica : ");
		m_lexica = new ObjectIntOpenHashMap<String>();
		
		for (i_N=0; (tree = reader.nextTree()) != null; i_N++)
		{
			for (String key : getLexica(tree))
			{
				if (!m_lexica.containsKey(key))
					m_lexica.put(key, d++);
			}
		}
		
		reader.close();
		System.out.println((i_D = m_lexica.size()));
	}
	
	private void initVectors(String filename)
	{
		DepReader reader = new DepReader(filename, true);
		DepTree   tree;
		
		System.out.print("Initializing vectors: ");
		a_vectors = new int[i_N][];
		
		for (int n=0; (tree = reader.nextTree()) != null; n++)
			a_vectors[n] = getVector(tree);
		
		reader.close();
		System.out.println(i_N);
	}
	
	private void initDistributions()
	{
		System.out.print("Initializing priors : ");
		initPriors();
		System.out.println(i_K+" x "+i_N);
		
		System.out.print("Initializing means  : ");
		initMeans();
		System.out.println(i_K);
		
		System.out.print("Initializing stdevs : ");
		initStdevs();
		System.out.println(i_K);
	}
	
	private void initPriors()
	{
		double p = 1d / i_K;
		
		p_k  = new double[i_K];
		p_kn = new double[i_N][i_K];
		
		Arrays.fill(p_k, p);
	}
	
	private void initMeans()
	{
		IntOpenHashSet set  = new IntOpenHashSet();
		Random         rand = new Random(0);
		int[]          vector;
		int            n, k = 0;
		
		d_m   = new double[i_D][i_K];
		d_mScalars = new double[i_K];
		
		while (set.size() < i_K)
		{
			n = rand.nextInt(i_N);
			
			if (!set.contains(n))
			{
				vector = a_vectors[n];
				set.add(n);
				
				for (int d : vector)
					d_m[d][k] = 1;

				d_mScalars[k++] = Math.sqrt(vector.length);
			}
		}
	}
	
	private void initStdevs()
	{
		double[] dist;
		int      n, k;
		d_s = new double[i_K];
		
		for (n=0; n<i_N; n++)
		{
			dist = getCosineDistances(a_vectors[n]);
			
			for (k=0; k<i_K; k++)
				d_s[k] += dist[k];
		}
		
		for (k=0; k<i_K; k++)
			d_s[k] = Math.sqrt(d_s[k] / i_N);
	}

	private HashSet<String> getLexica(DepTree tree)
	{
		HashSet<String> set = new HashSet<String>();
		
		addDepLexica(tree, set);
		return set;
	}
	
	protected void addDepLexica(DepTree tree, HashSet<String> set)
	{
		DepNode node, head;
		
		for (int i=1; i<tree.size(); i++)
		{
			node = tree.get(i);
			if (node.headId < 0)	continue;
			
			addDep1gramLexica(set, node, "<", node.deprel);
			if (node.headId == 0)	continue;
			head = tree.get(node.headId);
			addDep1gramLexica(set, head, ">", node.deprel);
			
			if (node.id < head.id)
			{
				addDep2gramLexica(set, node, head, "<", "");
				addDep2gramLexica(set, node, head, "<", node.deprel);
			}
			else
			{
				addDep2gramLexica(set, head, node, ">", "");
				addDep2gramLexica(set, head, node, ">", node.deprel);
			}
			
			if (head.headId < 0)	continue;
			addDep3gramLexica(set, node, head, tree.get(head.headId));
		}
	}
	
	private void addDep1gramLexica(HashSet<String> set, DepNode node, String dir, String deprel)
	{
		set.add(node.lemma + dir + deprel);
		set.add(node.pos   + dir + deprel);
	}
	
	private void addDep2gramLexica(HashSet<String> set, DepNode prev, DepNode next, String dir, String deprel)
	{
		String label = dir + deprel;
		
		set.add(prev.lemma + "_" + next.lemma + label);
		set.add(prev.lemma + "_" + next.pos   + label);
		set.add(prev.pos   + "_" + next.lemma + label);
		set.add(prev.pos   + "_" + next.pos   + label);
	}
	
	private void addDep3gramLexica(HashSet<String> set, DepNode node, DepNode head, DepNode grandHead)
	{
		set.add(node.pos+"_"+head.pos+"_"+grandHead.pos);
	}
	
	private int[] getVector(DepTree tree)
	{
		IntOpenHashSet set = new IntOpenHashSet();
		int d;
		
		for (String key : getLexica(tree))
		{
			if ((d = m_lexica.get(key)) > 0)
				set.add(d-1);
		}
		
		int[] vector = set.toArray();
		Arrays.sort(vector);
		
		return vector;
	}
	
	public void cluster()
	{
		double prevScore = 0, currScore = expectation();
		System.out.println("- Score : "+currScore);
		
		for (int i=0; i<MAX_ITER && Math.abs(currScore - prevScore) > STOP; i++)
		{
			System.out.println("\nIteration: "+i);
			prevScore = currScore;
			
			System.out.println("- M-step");
			maximization();
			
			System.out.println("- E-step");
			currScore = expectation();
			
			System.out.println("- Score : "+currScore);
		}
		
		Random rand = new Random();
		
		for (int i=0; i<10; i++)
			System.out.println(Arrays.toString(p_kn[rand.nextInt(i_N)]));
	}
	
	private double expectation()
	{
		double[] memberships;
		double   sum, score = 0;
		int      k, n;
		 
		for (n=0; n<i_N; n++)
		{
			memberships = getMemberships(a_vectors[n]);
			sum = 0;

			for (k=0; k<i_K; k++)
				sum += memberships[k];

			for (k=0; k<i_K; k++)
				p_kn[n][k] = memberships[k] / sum;
			
			score += Math.log(sum);
		}
		
		return score;
	}
	
	private double[] getMemberships(int[] vector)
	{
		double[] memberships = new double[i_K];
		double[] dist = getCosineDistances(vector);
		double   stdev;
		
		for (int k=0; k<i_K; k++)
		{
			stdev = d_s[k];
			memberships[k] = p_k[k] * (1 / (SQRT_2PI*stdev)) * Math.exp(-0.5 * dist[k] * (stdev*stdev));
		}
		
		return memberships;
	}
	
	private void maximization()
	{
		maximizePriors();
		maximizeMeans();
		maximizeStdDevs();
		
		for (int k=0; k<i_K; k++)
			p_k[k] /= i_N;
	}
	
	private void maximizePriors()
	{
		int n, k;
		
		Arrays.fill(p_k, 0);
		
		for (n=0; n<i_N; n++)
			for (k=0; k<i_K; k++)
				p_k[k] += p_kn[n][k];
	}
	
	private void maximizeMeans()
	{
		int n, k, d;
		double m;
		
		for (d=0; d<i_D; d++)
			Arrays.fill(d_m[d], 0);
		
		for (n=0; n<i_N; n++)
		{
			for (int idx : a_vectors[n])
				for (k=0; k<i_K; k++)
					d_m[idx][k] += p_kn[n][k];
		}
		
		Arrays.fill(d_mScalars, 0);
		
		for (d=0; d<i_D; d++)
		{
			for (k=0; k<i_K; k++)
			{
				m = d_m[d][k] / p_k[k];
				d_m[d][k] = m;
				d_mScalars[k] += m * m;
			}
		}
		
		for (k=0; k<i_K; k++)
			d_mScalars[k] = Math.sqrt(d_mScalars[k]);
	}
	
	private void maximizeStdDevs()
	{
		double[] dist;
		int n, k;
		
		Arrays.fill(d_s, 0);
		
		for (n=0; n<i_N; n++)
		{
			dist = getCosineDistances(a_vectors[n]);
			
			for (k=0; k<i_K; k++)
				d_s[k] += p_kn[n][k] * dist[k];
		}
		
		for (k=0; k<i_K; k++)
			d_s[k] = Math.sqrt(d_s[k] / p_k[k]);
	}
	
	private double[] getCosineDistances(int[] vector)
	{
		double[] dots = new double[i_K];
		double[] mean;
		double   scalar = Math.sqrt(vector.length);
		int      k;
		
		for (int d : vector)
		{
			mean = d_m[d];
			
			for (k=0; k<i_K; k++)
				dots[k] += mean[k];
		}
		
		for (k=0; k<i_K; k++)
			dots[k] = 1 - (dots[k] / (d_mScalars[k] * scalar));
		
		return dots;
	}

	static public void main(String[] args)
	{
		String trainFile = args[0];
		int    K = Integer.parseInt(args[1]);
		
		new DepGMCluster(trainFile, K);
	}
}
